from playLA.Vector import Vector

class Matrix:
    def __init__(self, list2d):
        # 每个向量长度应该相等
        # 元素类型只能是数字
        self._values = [row[:] for row in list2d]

    def __repr__(self):
        return "Matrix: {}".format(self._values)

    __str__ = __repr__

    def shape(self):
        # 形状： 行，列
        return len(self._values), len(self._values[0])

    def row_num(self):
        return self.shape()[0]

    __len__ = row_num

    def __getitem__(self, pos):
        r, c = pos
        return self._values[r][c]

    def col_num(self):
        return self.shape()[1]

    def size(self):
        r, c = self.row_num(), self.col_num()
        return r * c

    def row_vector(self, index):
        return Vector(self._values[index])

    def dot(self, other):
        if isinstance(other, Vector):
            assert self.col_num() == len(other), "Error in Matrix-Vector Multiplication"
            return Vector([self.row_vector(i).dot(other) for i in range(self.row_num())])
        if isinstance(other, Matrix):
            assert self.col_num() == other.row_num(), "Error in Matrix-Matrix Multiplication"
            return Matrix([[self.row_vector(i).dot(other.col_vector(j))
                            for j in range(other.col_num())] for i in range(self.row_num())])

    def T(self):
        return Matrix([[e for e in self.col_vector(i)]
                       for i in range(self.col_num())])

    def col_vector(self, index):
        return Vector([row[index] for row in self._values])

    def __add__(self, other):
        assert self.shape() == other.shape(), "error, shape mismatch"
        return Matrix([[a + b for a, b in zip(self.row_vector(i), other.row_vector(i))] for i in range(self.row_num())])

    def __sub__(self, other):
        assert self.shape() == other.shape(), "error, shape mismatch"
        return Matrix([[a - b for a, b in zip(self.row_vector(i), other.row_vector(i))] for i in range(self.row_num())])

    def __mul__(self, k):
        return Matrix([[e * k for e in self.row_vector(i)] for i in range(self.row_num())])

    def rmul(self, k):
        return self * k

    def __truediv__(self, k):
        return self * (1 / k)

    def __pos__(self):
        return self * 1

    def __neg__(self):
        return self * -1

    @classmethod
    def zeros(cls, row, col):
        return cls([[0] * col for _ in range(row)])

    @classmethod
    def identity(cls, n):
        """ 创建一个单位矩阵 n方阵 """
        m = [[0]*n for _ in range(n)]
        for i in range(n):
            m[i][i] = 1
        return cls(m)




